package trace

import (
	"context"
	"net/url"

	"gitee.com/AlexWoo/arpc"
	"go.opentelemetry.io/otel"
	"go.opentelemetry.io/otel/attribute"
	"go.opentelemetry.io/otel/exporters/jaeger"
	"go.opentelemetry.io/otel/propagation"
	"go.opentelemetry.io/otel/sdk/resource"
	sdktrace "go.opentelemetry.io/otel/sdk/trace"
	semconv "go.opentelemetry.io/otel/semconv/v1.21.0"
)

var (
	tr *Trace
)

const (
	traceLibname = "arpc"

	protocolUdp = "udp"
)

type Trace struct {
	conf     *TraceConf
	tp       *sdktrace.TracerProvider
	cluster  string
	instance string
}

func setupPropagator() {
	prop := propagation.NewCompositeTextMapPropagator(
		propagation.TraceContext{},
		propagation.Baggage{},
	)
	otel.SetTextMapPropagator(prop)
}

func newExporter(conf *TraceConf) (sdktrace.SpanExporter, error) {
	if u, err := url.Parse(conf.Endpoint); err == nil && u.Scheme == protocolUdp {
		return jaeger.New(jaeger.WithAgentEndpoint(jaeger.WithAgentHost(u.Hostname()),
			jaeger.WithAgentPort(u.Port())))
	} else {
		return jaeger.New(jaeger.WithCollectorEndpoint(jaeger.WithEndpoint(conf.Endpoint)))
	}
}

func setupTraceProvider(conf *TraceConf, cluster string, instance string) (*sdktrace.TracerProvider, error) {
	exp, err := newExporter(conf)
	if err != nil {
		return nil, err
	}

	attr := []attribute.KeyValue{
		semconv.ServiceNameKey.String(conf.ServiceName),
	}

	if cluster != "" {
		attr = append(attr, semconv.K8SClusterName(cluster))
	}

	if instance != "" {
		attr = append(attr, semconv.K8SContainerName(instance))
	}

	opts := []sdktrace.TracerProviderOption{
		sdktrace.WithSampler(sdktrace.ParentBased(sdktrace.TraceIDRatioBased(conf.Sampler))),
		sdktrace.WithResource(resource.NewSchemaless(attr...)),
		sdktrace.WithBatcher(exp),
	}

	provider := sdktrace.NewTracerProvider(opts...)
	otel.SetTracerProvider(provider)

	return provider, nil
}

func CreateTrace(conf *TraceConf, cluster string, instance string) error {
	conf.ReplaceEnv()

	if conf.Endpoint == "" {
		return nil
	}

	if tr != nil {
		return nil
	}

	// Set up propagator
	setupPropagator()

	// Set up trace provider
	tp, err := setupTraceProvider(conf, cluster, instance)
	if err != nil {
		return err
	}

	otel.SetErrorHandler(otel.ErrorHandlerFunc(func(err error) {
		arpc.Errorf("otel error: %v", err)
	}))

	tr = &Trace{
		conf:     conf,
		tp:       tp,
		cluster:  cluster,
		instance: instance,
	}

	return nil
}

func StopTrace() {
	if tr == nil {
		return
	}

	tr.tp.Shutdown(context.Background())
}
